﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading;
using Microsoft.Win32.SafeHandles;
using Mono.Cecil.PE;
using NHook.Native;

namespace NHook
{

	public class PatchesCollection : List<Patch>
	{
		ProcessDebugger _Debugger;
		public PatchesCollection(ProcessDebugger debugger)
		{
			_Debugger = debugger;
		}

		public Patch New(VA instruction)
		{
			var patch = new Patch(_Debugger.Debuggee.Information.ProcessHandle, instruction);
			this.Add(patch);
			return patch;
		}

		public void ModifyImage(ProcessModule processModule, string patchedExecutable)
		{
			File.Copy(processModule.FileName, patchedExecutable, true);
			using(var fs = File.Open(patchedExecutable, FileMode.Open))
			{
				foreach(var patch in this)
				{
					patch.ApplyPatch(processModule, fs);
				}
			}
		}
	}
	public class BreakpointCollection : Dictionary<VA, Breakpoint>
	{
		ProcessDebugger _Debugger;
		public BreakpointCollection(ProcessDebugger debugger)
		{
			_Debugger = debugger;
		}


		public Breakpoint AtEntryPoint()
		{
			var process = _Debugger.EnsureProcessLoaded();
			var mainImage = Image.ReadFromFile(process.MainModule.FileName);
			var instructionAddress = new VA(process.MainModule.BaseAddress) + new RVA(mainImage.NtHeader.EntryPoint);
			var breakpoint = NewBreakpoint(instructionAddress);
			return breakpoint;
		}

		private Breakpoint NewBreakpoint(VA instructionAddress)
		{
			_Debugger.EnsureProcessLoaded();
			var breakpoint = new Breakpoint(_Debugger, instructionAddress);
			breakpoint.Toggle();
			this.Add(instructionAddress, breakpoint);
			return breakpoint;
		}

		public Breakpoint At(string module, RVA rva)
		{
			_Debugger.EnsureProcessLoaded();
			var moduleAddress = _Debugger.AddressOfModule(module);
			return NewBreakpoint(moduleAddress + rva);
		}
	}

	public class Debuggee
	{
		ProcessDebugger _Debugger;
		public Debuggee(ProcessDebugger debugger, ProcessInformation processInfo)
		{
			Information = processInfo;
			this._Debugger = debugger;
		}
		public ProcessInformation Information
		{
			get;
			private set;
		}

		Process _Process;
		internal void SetProcess()
		{
			_Process = Process.GetProcessById((int)Information.ProcessId);
		}
		public Process Process
		{
			get
			{
				return _Process;
			}
		}

		public ProcessModule MainModule
		{
			get
			{
				if(Process == null)
					return null;
				return Process.MainModule;
			}
		}

		public IList<ObjectHandle> GetHandles()
		{
			_Debugger.EnsureProcessLoaded();
			return ObjectHandle.GetObjectHandles(Information);
		}
	}
	public class ProcessDebugger : IDisposable, Commandable
	{
		internal enum WaitBehavior
		{
			Keep,
			Skip
		}


		private readonly Dictionary<uint, ThreadDebugger> _ThreadDebuggers = new Dictionary<uint, ThreadDebugger>();
		public Dictionary<uint, ThreadDebugger> ThreadDebuggers
		{
			get
			{
				return _ThreadDebuggers;
			}
		}

		private PatchesCollection _Patches;
		public PatchesCollection Patches
		{
			get
			{
				if(_Patches == null)
					_Patches = new PatchesCollection(this);
				return _Patches;
			}
		}

		Debuggee _Debuggee;
		public Debuggee Debuggee
		{
			get
			{
				return _Debuggee;
			}
			private set
			{
				_Debuggee = value;
			}
		}

		DebugEvent _CurrentEvent;


		public DebugEvent CurrentEvent
		{
			get
			{
				return _CurrentEvent;
			}
		}
		public ThreadDebugger BreakingThread
		{
			get
			{
				if(_CurrentEvent == null)
					return null;
				return ThreadDebuggers[_CurrentEvent.ThreadId];
			}
		}


		private void AssertIsBreaking()
		{
			if(_CurrentEvent == null)
				throw new InvalidOperationException("The debugger is not breaking");
		}

		public bool IsStarted
		{
			get;
			private set;
		}
		public ProcessDebugger()
		{
			_Wait = new Wait(this);
			_Continue = new Continue(this);
		}
		public ProcessInformation Start(string path)
		{
			if(IsStarted)
				throw new InvalidOperationException("Debugger already started");
			IsStarted = true;
			System.Diagnostics.Process.EnterDebugMode();
			var processInfo = NHook.Native.Helper.StartDebugProcess(path);
			Debuggee = new Debuggee(this,processInfo);
			_SymbolManager = new SymbolManager(this);
			return processInfo;
		}

		public ProcessInformation Attach(int pid)
		{
			if(IsStarted)
				throw new InvalidOperationException("Debugger already started");
			IsStarted = true;
			System.Diagnostics.Process.EnterDebugMode();
			var processInfo = NHook.Native.Helper.DebugActiveProcess(pid);
			Debuggee = new Debuggee(this,processInfo);
			_SymbolManager = new SymbolManager(this);
			return processInfo;
		}

		BreakpointCollection _Breakpoints;
		public BreakpointCollection Breakpoints
		{
			get
			{
				if(_Breakpoints == null)
					_Breakpoints = new BreakpointCollection(this);
				return _Breakpoints;
			}
		}

		public void Detach()
		{
			Stop(true);
		}

		void Stop(bool detach)
		{
			if(!IsStarted)
				throw new InvalidOperationException("Debugger not started");

			if(!Debuggee.Process.HasExited)
			{
				foreach(var breakpoint in Breakpoints)
				{
					breakpoint.Value.Disable();
				}

				foreach(var thread in ThreadDebuggers.Values)
					thread.StepByStep = false;

				if(Debuggee.Process != null && !Debuggee.Process.HasExited)
					Continue.UntilNoBreakpointOrSingleStep();
				if(detach)
					NHook.Native.Helper.DebugSetProcessKillOnExit(!detach);
				NHook.Native.Helper.DebugActiveProcessStop(Debuggee.Information);
			}
			IsStarted = false;
			Debuggee = null;
			System.Diagnostics.Process.LeaveDebugMode();
		}
		public void Stop()
		{
			Stop(false);
		}



		readonly Continue _Continue;
		public Continue Continue
		{
			get
			{
				return _Continue;
			}
		}
		readonly Wait _Wait;
		public Wait Wait
		{
			get
			{
				return _Wait;
			}
		}

		public bool IsBreaking
		{
			get;
			private set;
		}
		internal DebugEvent WaitNextEvent(out WaitBehavior waitBehavior)
		{
			AssertStarted();
			if(IsBreaking)
				_Continue.Run(true);
			if(IsBreaking)
				throw new InvalidOperationException("NHook BUG : IsBreaking should be false now");

			DebugEvent dbgEvent = NHook.Native.Helper.WaitForEvent(Debuggee.Information, TimeSpan.MaxValue);
			_CurrentEvent = dbgEvent;
			IsBreaking = true;

			if(dbgEvent.EventType == DebugEventType.CreateThreadEvent)
			{
				ThreadDebuggers.Add(dbgEvent.ThreadId, new ThreadDebugger(dbgEvent.As<CreateThreadDetail>().ThreadHandle, this, dbgEvent.ThreadId));
			}
			if(dbgEvent.EventType == DebugEventType.CreateProcessEvent)
			{
				ThreadDebuggers.Add(dbgEvent.ThreadId, new ThreadDebugger(Debuggee.Information.ThreadHandle, this, dbgEvent.ThreadId));
			}
			if(dbgEvent.EventType == DebugEventType.ExitProcessEvent)
			{
				ExitCode = dbgEvent.As<ExitProcessDetail>().ExitCode;
			}
			if(Debuggee.Process == null)
			{
				if(dbgEvent.EventType == DebugEventType.LoadDllEvent)
				{
					if(AreSameImage(dbgEvent.As<LoadDllDetail>().ImageName, "kernel32.dll"))
					{
						Debuggee.SetProcess();
						if(!SymbolManager.ManualLoading)
							SymbolManager.LoadSymbolsFromExe(Debuggee.MainModule.FileName);
					}
					if(!SymbolManager.ManualLoading)
						SymbolManager.LoadSymbolsFromExe(dbgEvent.As<LoadDllDetail>().Location);
				}
			}
			BreakingThread.CurrentEvent = dbgEvent;
			waitBehavior = BreakingThread.Wait.FireOnWaitNext(dbgEvent);
			Wait.FireOnWaitNext(dbgEvent);
			return dbgEvent;
		}




		private void AssertStarted()
		{
			if(!IsStarted)
				throw new InvalidOperationException("Debugger not running, call Start(...) first");
		}

		public void Dispose()
		{
			if(IsStarted)
				Stop();
		}
		~ProcessDebugger()
		{
			if(IsStarted)
				throw new InvalidOperationException("A debugger is leaking, you have to Stop or Dispose them on the thread that created it.");

		}

		internal void WriteMemory(VA address, byte instruction)
		{
			var bytes = new byte[] { instruction };
			NHook.Native.MemoryUtil.WriteMemory(Debuggee.Information.ProcessHandle.DangerousGetHandle(), new IntPtr(address.Address), bytes);
		}

		internal byte ReadMemory(VA address)
		{
			return NHook.Native.MemoryUtil.ReadMemory(Debuggee.Information.ProcessHandle.DangerousGetHandle(), new IntPtr(address.Address), 1)[0];
		}

		public int? ExitCode
		{
			get;
			set;
		}

		internal static bool AreSameImage(string path1, string path2)
		{
			if(String.IsNullOrEmpty(path1))
				return false;
			return Path.GetFileName(path1) == path2;
		}


		internal void ContinueNextEvent(DebugEvent evt, bool handleException = true)
		{
			AssertStarted();
			var threadDebugger = ThreadDebuggers[evt.ThreadId];
			threadDebugger.Continue.FireContinuing();
			Continue.FireContinuing();
			NHook.Native.Helper.Continue(evt, handleException);
			IsBreaking = false;
			_CurrentEvent = null;
			ThreadDebuggers[evt.ThreadId].CurrentEvent = null;
		}

		#region Commandable Members

		Wait Commandable.Wait
		{
			get
			{
				return Wait;
			}
		}

		Continue Commandable.Continue
		{
			get
			{
				return Continue;
			}
		}

		ProcessDebugger Commandable.ProcessDebugger
		{
			get
			{
				return this;
			}
		}

		ThreadDebugger Commandable.GetCurrentThreadDebugger()
		{
			return BreakingThread;
		}

		DebugEvent Commandable.CurrentEvent
		{
			get
			{
				return _CurrentEvent;
			}
		}

		#endregion

		private SymbolManager _SymbolManager;
		public SymbolManager SymbolManager
		{
			get
			{
				return _SymbolManager;
			}
		}

		public VA AddressOfModule(string moduleName)
		{
			if(moduleName == null)
				throw new ArgumentNullException("moduleName");
			Debuggee.Process.Refresh();
			var module = Debuggee.Process.Modules.OfType<ProcessModule>().FirstOrDefault(m => String.Equals(m.ModuleName, moduleName, StringComparison.InvariantCultureIgnoreCase));
			if(module == null)
				throw new InvalidOperationException("Module " + moduleName + " not loaded into process");
			return new VA(module.BaseAddress);
		}




		internal System.Diagnostics.Process EnsureProcessLoaded()
		{
			var process = Debuggee.Process;
			if(process == null)
			{
				Continue.UntilDllLoaded("kernel32.dll");
				process = Debuggee.Process;
			}
			return process;
		}

	}
}
